"""
Variational Data Assimilation Framework (VDAF)
Complete implementation with examples for both 3D-Var and 4D-Var using barotropic model
"""
import numpy as np
from scipy.optimize import minimize
from typing import Dict, List, Optional, Tuple
import matplotlib.pyplot as plt
import random
from Tools.AssimilationSystem import *
from Tools.var_methods import *
from Tools.adapters import *
from Tools.var4d_capture import create_var4d_with_capture
from Tools.barotropic_model import BarotropicModel
from Tools.barotropic_adapter import BarotropicObservationOperator
import matplotlib.pyplot as plt
# import cupy as cp
import time  # Add time module for measuring execution time

def create_assimilation_config(method, state_size, obs_size, time_window, bg_error_variance, obs_error_variance):
    """Create and return an assimilation configuration."""
    return AssimilationConfig(
        method=method,
        state_size=state_size,
        obs_size=obs_size,
        time_window=time_window,
        max_outer_loops=1,
        max_inner_loops=100,
        bg_error_variance=bg_error_variance,
        obs_error_variance=obs_error_variance
    )

def run_assimilation(method, model, x_b, observations, obs_ops, state_size, obs_size, time_window, 
                     bg_error_variance, obs_error_variance):
    """Run the data assimilation and return analysis and timing."""
    config = create_assimilation_config(
        method, state_size, obs_size, time_window, bg_error_variance, obs_error_variance
    )
    
    assim_system = AssimilationSystem(config, model)
    assim_system.set_observation_operators(obs_ops)
    
    start_time = time.time()
    x_a = assim_system.assimilate(x_b, observations)
    assim_time = time.time() - start_time
    
    return x_a, assim_time

def calculate_rmse(x, x_truth):
    """Calculate RMSE between state and truth."""
    return np.sqrt(np.mean((x - x_truth)**2))

def integrate_trajectories(model, states, time_window):
    """Integrate multiple initial states forward and return trajectories and timing."""
    start_time = time.time()
    trajectories = {name: model.integrate(state, time_window) for name, state in states.items()}
    integration_time = time.time() - start_time
    return trajectories, integration_time

def plot_state_comparison(model, x_truth, x_b, x_a, title_suffix="", filename="comparison_results.png"):
    """Plot comparison between truth, background and analysis state."""
    plt.figure(figsize=(15, 5))
    
    # Plot vorticity fields
    plt.subplot(1, 3, 1)
    plt.title(f"Truth {title_suffix}")
    truth_2d = model._reshape_state(x_truth)
    plt.imshow(truth_2d, cmap='RdBu_r')
    plt.colorbar(label='Vorticity')
    
    plt.subplot(1, 3, 2)
    plt.title(f"Background {title_suffix}")
    bg_2d = model._reshape_state(x_b)
    plt.imshow(bg_2d, cmap='RdBu_r')
    plt.colorbar(label='Vorticity')
    
    plt.subplot(1, 3, 3)
    plt.title(f"Analysis {title_suffix}")
    analysis_2d = model._reshape_state(x_a)
    plt.imshow(analysis_2d, cmap='RdBu_r')
    plt.colorbar(label='Vorticity')
    
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

def plot_time_evolution(model, trajectories, time_window, filename="time_evolution.png"):
    """Plot the time evolution of different trajectories."""
    n_trajectories = len(trajectories)
    cols = 3  # Fixed number of columns
    rows = 2  # Two rows: initial and final states
    

    
    for i, (name, trajectory) in enumerate(trajectories.items()):
        # Initial state (top row)
        plt.subplot(rows, cols, i + 1)
        plt.title(f"{name} Initial (t=0)")
        plt.imshow(model._reshape_state(trajectory[0]), cmap='RdBu_r')
        plt.colorbar(label='Vorticity')
        
        # Final state (bottom row)
        plt.subplot(rows, cols, i + 1 + cols)
        plt.title(f"{name} Final (t={time_window})")
        plt.imshow(model._reshape_state(trajectory[time_window]), cmap='RdBu_r')
        plt.colorbar(label='Vorticity')
    
    plt.tight_layout()
    plt.savefig(filename)
    plt.close()

def print_results(method, rmse_b, rmse_a, assim_time):
    """Print the results of data assimilation."""
    print(f"\nResults for {method} with Barotropic Model:")
    print(f"Background RMSE: {rmse_b:.4f}")
    print(f"Analysis RMSE: {rmse_a:.4f}")
    print(f"Improvement: {(1 - rmse_a/rmse_b)*100:.1f}%")
    print(f"{method} assimilation time: {assim_time:.4f} seconds")

def print_performance_metrics(integration_time, forecast_time, assim_times):
    """Print performance metrics."""
    print("\n=== Performance Metrics ===")
    print(f"Truth trajectory integration time: {integration_time:.4f} seconds")
    print(f"Background & analysis trajectory time: {forecast_time:.4f} seconds")
    
    for method, time_value in assim_times.items():
        print(f"{method} assimilation time: {time_value:.4f} seconds")
    
    # Print ratios if there are multiple methods
    methods = list(assim_times.keys())
    if len(methods) > 1:
        for i in range(len(methods)):
            for j in range(i+1, len(methods)):
                ratio = assim_times[methods[i]] / assim_times[methods[j]]
                print(f"{methods[i]}/{methods[j]} time ratio: {ratio:.2f}x")

def run_3dvar_example(model, x_b, x_truth, all_observations, all_obs_ops, 
                      state_size, obs_size, bg_error_variance, obs_error_variance):
    """Run the 3D-Var assimilation example."""
    print("\n=== Running 3D-Var Example with Barotropic Model ===")
    
    # Create 3D-Var system with t=0 observations only
    observations_3dvar = {0: all_observations[0]}
    x_a_3dvar, var3d_time = run_assimilation(
        '3DVar', model, x_b, observations_3dvar, [all_obs_ops[0]], 
        state_size, obs_size, 0, bg_error_variance, obs_error_variance
    )
    
    # Calculate error statistics
    rmse_b = calculate_rmse(x_b, x_truth)
    rmse_a_3dvar = calculate_rmse(x_a_3dvar, x_truth)
    
    # Print results
    print_results('3D-Var', rmse_b, rmse_a_3dvar, var3d_time)
    
    # Visualize results
    plot_state_comparison(model, x_truth, x_b, x_a_3dvar, "", "barotropic_3dvar_results.png")
    
    return x_a_3dvar, var3d_time, rmse_a_3dvar

def run_4dvar_example(model, x_b, x_truth, all_observations, all_obs_ops, 
                      state_size, obs_size, max_time_window, bg_error_variance, obs_error_variance):
    """Run the 4D-Var assimilation example."""
    print("\n=== Running 4D-Var Example with Barotropic Model ===")
    
    # Run 4D-Var with the barotropic model
    x_a_4dvar, var4d_time = run_assimilation(
        '4DVar', model, x_b, all_observations, all_obs_ops,
        state_size, obs_size, max_time_window, bg_error_variance, obs_error_variance
    )
    
    # Calculate error statistics
    rmse_b = calculate_rmse(x_b, x_truth)
    rmse_a_4dvar = calculate_rmse(x_a_4dvar, x_truth)
    
    # Print results
    print_results('4D-Var', rmse_b, rmse_a_4dvar, var4d_time)
    
    # Visualize results
    plot_state_comparison(model, x_truth, x_b, x_a_4dvar, "", "barotropic_4dvar_results.png")
    
    return x_a_4dvar, var4d_time, rmse_a_4dvar


def run_drp4dvar(x_b, x_truth, model, state_size, obs_size, max_time_window, bg_error_variance, 
                obs_error_variance, all_obs_ops, all_observations):
    """Run DRP4DVar assimilation with observations from all time steps."""
    print("\n=== Running DRP4DVar Example ===")
    config_4dvar = AssimilationConfig(
        method='DRP4DVar',
        state_size=state_size,
        obs_size=obs_size,
        time_window=max_time_window,
        max_outer_loops=1,
        max_inner_loops=50,
        bg_error_variance=bg_error_variance,
        obs_error_variance=obs_error_variance
    )
    config_4dvar.ensemble_size = 30
    
    # Run DRP4DVar with the same initial conditions
    assim_system_4dvar = AssimilationSystem(config_4dvar, model)
    assim_system_4dvar.set_observation_operators(all_obs_ops)
    x_a_4dvar = assim_system_4dvar.assimilate(x_b, all_observations)
    
    # Calculate DRP4DVar error statistics
    rmse_b = np.sqrt(np.mean((x_b - x_truth)**2))
    rmse_a_4dvar = np.sqrt(np.mean((x_a_4dvar - x_truth)**2))
    
    print(f"\nResults for DRP4DVar:")
    # print(f"max inner loops: {config_4dvar.max_inner_loops}")
    #print(f"ensemble size: {config_4dvar.ensemble_size}")
    print(f"Background RMSE: {rmse_b:.4f}")
    print(f"Analysis RMSE: {rmse_a_4dvar:.4f}")
    print(f"Improvement: {(1 - rmse_a_4dvar/rmse_b)*100:.1f}%")
    plot_state_comparison(model, x_truth, x_b, x_a_4dvar, "", "barotropic_drp4dvar_results.png")
    
    return x_a_4dvar, rmse_b, rmse_a_4dvar


def main():
    # Set fixed seeds for reproducibility
    RANDOM_SEED = 42
    random.seed(RANDOM_SEED)
    np.random.seed(RANDOM_SEED)
    
    # Create consistent configuration
    nx, ny = 10, 10  # Grid dimensions
    state_size = nx * ny  # State vector size
    obs_size = state_size // 4  # Number of observations (25% of state)
    bg_error_variance = 0.5
    obs_error_variance = 0.1
    
    # Create barotropic model instead of the simple linear model
    model = BarotropicModel(nx=nx, ny=ny)

    # For NVIDIA Hopper (SM_89), set device optimization
    # if cp.cuda.runtime.getDeviceCount() > 0:
    #     cp.cuda.Device(0).use()  # Use first GPU
    #     # Optional: Set kernel optimization
    #     cp.cuda.set_allocator(cp.cuda.MemoryPool().malloc)
    #     # Print GPU info
    #     print(f"Using GPU: {cp.cuda.runtime.getDeviceProperties(0)['name'].decode()}")
    
    # Create observation operators for all potential time steps
    max_time_window = 4  # Maximum time window for 4D-Var
    all_obs_ops = []
    for t in range(max_time_window + 1):
        # Use same random seed to ensure consistent placement
        random.seed(42 + t)
        # Use the specialized barotropic observation operator
        op = BarotropicObservationOperator(model, state_size, obs_size, obs_type='vorticity')
        all_obs_ops.append(op)
    
    # Generate synthetic truth state
    x_truth = model.generate_random_state(mean_amplitude=1.0, seed=RANDOM_SEED)
    
    # Generate one set of synthetic data with the 4D-Var time window
    config_data = AssimilationConfig(
        method='DRP4DVar',
        state_size=state_size,
        obs_size=obs_size,
        time_window=max_time_window,
        bg_error_variance=bg_error_variance,
        obs_error_variance=obs_error_variance
    )
    
    # Generate background state by perturbing truth
    error = np.random.multivariate_normal(np.zeros(state_size), config_data.b_matrix)
    x_b = x_truth + error
    
    # Generate observations by running truth model
    start_time = time.time()
    truth_traj = model.integrate(x_truth, max_time_window)
    integration_time = time.time() - start_time
    print(f"\nModel integration time (truth trajectory): {integration_time:.4f} seconds")
    
    all_observations = {}
    for t in range(max_time_window + 1):
        # Apply observation operator
        y_perfect = all_obs_ops[t].h(truth_traj[t])
        
        # Add observation error
        obs_error = np.random.multivariate_normal(
            np.zeros(obs_size), 
            config_data.r_matrix
        )
        y_obs = y_perfect + obs_error
        
        all_observations[t] = y_obs
    
    # Calculate background RMSE 
    rmse_b = calculate_rmse(x_b, x_truth)
    
    # Uncomment to run 3D-Var
    # x_a_3dvar, var3d_time, rmse_a_3dvar = run_3dvar_example(
    #     model, x_b, x_truth, all_observations, all_obs_ops,
    #     state_size, obs_size, bg_error_variance, obs_error_variance
    # )
    
    # # Run 4D-Var
    # start = time.perf_counter()
    # x_a_4dvar, var4d_time, rmse_a_4dvar = run_4dvar_example(
    #     model, x_b, x_truth, all_observations, all_obs_ops,
    #     state_size, obs_size, max_time_window, bg_error_variance, obs_error_variance
    # )
    # elapsed_4dvar = time.perf_counter() - start
    # print(f"run_4dvar time: {elapsed_4dvar:.3f} seconds")
    start = time.perf_counter()
    x_a_drp4dvar, _, rmse_a_drp4dvar = run_drp4dvar(x_b, x_truth, model, state_size, obs_size, 
                                                  max_time_window, bg_error_variance, 
                                                  obs_error_variance, all_obs_ops, all_observations)
    elapsed_drp4dvar = time.perf_counter() - start
    print(f"run_drp4dvar time: {elapsed_drp4dvar:.3f} seconds")
    
    
    # Integrate trajectories
    trajectories = {
        "Truth": x_truth,
        "Background": x_b,
        # "3D-Var Analysis": x_a_3dvar,  # Uncomment if running 3D-Var
        #"4D-Var Analysis": x_a_4dvar,
        "Analysis": x_a_drp4dvar
    }
    evolved_trajectories, forecast_time = integrate_trajectories(model, trajectories, max_time_window)
    
    # Print model timing statistics
    model.print_timing_statistics()
    
    # Plot time evolution
    plot_time_evolution(model, evolved_trajectories, max_time_window, "barotropic_time_evolution.png")
    
    # Print comparison of methods
    print("\n=== Comparison of Methods with Barotropic Model ===")
    # print(f"3D-Var RMSE: {rmse_a_3dvar:.4f}")  # Uncomment if running 3D-Var
    # print(f"4D-Var RMSE: {rmse_a_4dvar:.4f}")
    print(f"DRP4D-Var RMSE: {rmse_a_drp4dvar:.4f}")
    
    # Print performance metrics
    assim_times = {
        # "3D-Var": var3d_time,  # Uncomment if running 3D-Var
        # "4D-Var": var4d_time,
        "DRP4D-Var": elapsed_drp4dvar
    }
    print_performance_metrics(integration_time, forecast_time, assim_times)

if __name__ == "__main__":
    main()
